import cv2
import myfuncs as mf
import numpy as np
from matplotlib import pyplot as plt
import Template as tp
from ColorFilter import ColorFilter
from Kernels import Kernels


class MatchDigit:
	# kernel的大小要与原始图像的大小相匹配

	originImage = None
	grayImage = None
	binaryImage = None
	imageROI = None
	rectLocations = []
	singleTemplateSize = ()
	matchResults = []
	templateRectRatio = None
	tolerance = {'width': 20, 'height': 10, 'ratio': 0.12}
	kernels = Kernels()

	def __init__(self, imagePath, width, singleTemplateSize):
		self.imageInitialize(imagePath, width=width)
		self.setSingleTemplateSize(singleTemplateSize)
		self.templateRectRatio = singleTemplateSize[1] / (singleTemplateSize[0] * 4)

	def setSingleTemplateSize(self, singleTemplateSize):
		self.singleTemplateSize = singleTemplateSize

	def imageInitialize(self, imagePath, width):
		originImage = cv2.imread(imagePath)
		self.originImage = mf.resize(originImage, width=width)
		colorFilter = ColorFilter(self.originImage)
		colorFilter.addColor("orange", [11, 43, 46], [34, 255, 255])
		colorFilter.addColor("black", [0, 0, 0], [100, 100, 100])
		orangeArea = colorFilter.filterColor("orange")[1]
		blackArea = colorFilter.filterColor("black")[1]
		filteredImage = cv2.bitwise_or(orangeArea, blackArea)
		self.grayImage = cv2.cvtColor(self.originImage, cv2.COLOR_BGR2GRAY)
		self.binaryImage = filteredImage

	def calcROI(self):
		# cv2.morphologyEx()形态学操作通用函数
		imageROI = self.binaryImage.copy()
		# cv2.imshow("imageROI", imageROI)
		imageROI = cv2.morphologyEx(imageROI, cv2.MORPH_CLOSE, self.kernels.getKernel('9r5'))
		# cv2.imshow("1 close", imageROI)
		imageROI = cv2.morphologyEx(imageROI, cv2.MORPH_DILATE, self.kernels.getKernel('5r5'))
		# cv2.imshow("2 DILATE", imageROI)
		imageROI = cv2.morphologyEx(imageROI, cv2.MORPH_DILATE, self.kernels.getKernel('3r3'))
		# cv2.imshow("3 DILATE", imageROI)
		imageROI = cv2.morphologyEx(imageROI, cv2.MORPH_CLOSE, self.kernels.getKernel('9r1'))
		# cv2.imshow("4 CLOSE", imageROI)
		# mf.waitEnd()
		return imageROI

	def findContours(self, roi):
		contours, hierarchy = cv2.findContours(roi.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
		contoursImage = self.originImage.copy()
		cv2.drawContours(contoursImage, contours, -1, (200, 20, 20), 2)
		return contours, contoursImage

	def findRectROI(self, imageROI):
		contours, _ = self.findContours(imageROI)
		rectLocations = []
		# rectParams = []
		# 计算矩形，此处的矩形是一组的矩形
		for (i, contour) in enumerate(contours):
			(x, y, w, h) = cv2.boundingRect(contour)
			# 长宽比
			operandRatio = h / float(w)
			# 选择合适的区域
			# print("width: {: >5}| height: {: >5}| ratio: {: >5}".format(w, h, operandRatio))
			if abs(self.templateRectRatio - operandRatio) < self.tolerance['ratio']:
				minWidth = self.singleTemplateSize[0] * 5 - self.tolerance['width']
				maxWidth = self.singleTemplateSize[0] * 5 + self.tolerance['width']
				minHeight = self.singleTemplateSize[1] - self.tolerance['height']
				maxHeight = self.singleTemplateSize[1] + self.tolerance['height']
				# print("{0:-^20}".format("ratio matched"))
				if (minWidth <= w <= maxWidth) and (minHeight <= h <= maxHeight):
					# rectParams.append((x, y, w, h))
					# print("{0:-^20}".format("hit"))
					# 下面的操作通过构建一个边框，提取边框于灰度图中的颜色值，计算颜色值的平均值作为过滤条件
					marginColor = 0
					flag = 0
					mask = np.zeros([h + 2, w + 20], np.uint8)
					mask[0:1, 0:w + 20] = 1
					mask[-1:h + 2, 0:w + 20] = 1
					mask[0:h + 2, 0:3] = 1
					mask[0:h + 2, -3:w + 20] = 1
					if not (mask.shape == self.grayImage[y - 1:y + h + 1, x - 17:x + w + 3].shape):
						continue
					margin = cv2.bitwise_and(self.grayImage[y - 1:y + h + 1, x - 17:x + w + 3], self.grayImage[y - 1:y + h + 1, x - 17:x + w + 3], mask=mask)
					for r in range(0, len(margin)):
						for c in range(0, len(margin[r])):
							if margin[r, c] != 0:
								marginColor += margin[r, c]
								flag += 1
					marginColor /= flag
					# print(marginColor)
					if abs(marginColor - 255) < 0.1:
						rectLocations.append((x, y, w, h))
		# rectLocations = sorted(rectLocations, key=lambda x: x[0])
		self.rectLocations = rectLocations

		# r = mf.drawRects(rectParams, self.originImage.shape)
		# mf.showImage("r", r)

	def matchDigit(self, templates):
		imageROI = self.calcROI()
		self.findRectROI(imageROI)
		# templates is a dictionary of singleTemplates.
		# gX is the x position of a group
		matchResults = []
		for (i, (gX, gY, gW, gH)) in enumerate(self.rectLocations):
			# 此处的 + - 5相当于一个外边距
			group = self.grayImage[gY - 5:gY + gH + 5, gX - 5:gX + gW + 5]
			group = cv2.threshold(group, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)[1]
			text = ''
			groupOutput = {}
			# plt.subplot(111), plt.imshow(group), plt.show()
			# subContours: 4个数字的轮廓
			subContours, hierarchy = cv2.findContours(group.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
			subContours = mf.sortContours(subContours, 'left-to-right')[0]
			for subContour in subContours:
				(x, y, w, h) = cv2.boundingRect(subContour)
				roi = group[y:y+h, x:x+w]
				# plt.subplot(111), plt.title(roi.shape[0]/roi.shape[1]), plt.imshow(roi), plt.show()
				if roi.shape[0] < self.singleTemplateSize[0] * 0.3 and roi.shape[1] < self.singleTemplateSize[1] * 0.3:
					result = '12'
				elif abs(roi.shape[0] / roi.shape[1] - 0.2) < self.tolerance['ratio']:
					result = '11'
				else:
					roi = cv2.threshold(roi.copy(), 0, 255, cv2.THRESH_BINARY)[1]
					# mf.showImage('roi', roi)
					# dict.items() returns a enumerable tuple list.
					scores = []
					for (templateIndex, template) in templates.items():
						# cv2.matchTemplate() will return a 2-dimensional array that
						# concludes the match level of each point.
						result = cv2.matchTemplate(roi, template, cv2.TM_CCOEFF)
						# get the max value of match result
						(_, score, _, _) = cv2.minMaxLoc(result)
						scores.append(score)
					# np.argmax() returns the index of a array's maximum.
					# plt.subplot(121), plt.imshow(roi), plt.title('roi')
					# plt.subplot(122), plt.imshow(templates[np.argmax(scores)]), plt.title('matched')
					# plt.show()
					result = str(np.argmax(scores))

				if result == '10':
					text += '+'
				elif result == '11':
					text += '-'
				elif result == '12':
					text += '.'
				else:
					text += result
			groupOutput = {
				'id': i,
				'position': (gX, gY),
				'size': (gW, gH),
				'content': text
			}
			matchResults.append(groupOutput)
		self.matchResults = matchResults

	def drawResults(self):
		canvas = self.originImage.copy()
		# print(self.matchResults)
		for result in self.matchResults:
			text = result['content']
			x, y = result['position']
			w, h = result['size']
			cv2.rectangle(canvas, (x - 5, y - 5), (x + w + 5, y + h + 5), (200, 20, 20), 2)
			cv2.putText(canvas, text, (x - 5, y - 15), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (200, 20, 20), 2)
		return canvas

	def getMatchResults(self):
		return self.matchResults

	def getOperatedArea(self):
		w = self.originImage.shape[0]
		h = self.originImage.shape[1]
		hit = np.zeros([w, h], np.uint8)
		miss = np.ones([w, h], np.uint8)
		for result in self.matchResults:
			x, y = result['position']
			w, h = result['size']

			hit[y:y + h, x:x + w] = 1
			miss[y:y + h, x:x + w] = 0
		hit *= 255
		miss *= 255
		return hit, miss
